#include "C3DField.h"
#include <fftw3.h>


C3DField::C3DField() {
  Grid = 0;
  Width = -1;
  Min = CVector(1,1,1)*1E30;
  Max = CVector(1,1,1)*-1E30;
  PlanFFTInitialized=false;
  PlanFFTInvInitialized=false;
  in = 0;
  out = 0;
  noval = -1;
}


C3DField::~C3DField() {
  Delete();

  if (in) 
    fftw_free(in); 
  if (out)
    fftw_free(out);

  if (PlanFFTInitialized)
    fftw_destroy_plan(PlanFFT);

  if (PlanFFTInvInitialized)
    fftw_destroy_plan(PlanFFTInv);

}

void C3DField::Delete() {
  if (Grid)
    delete[] Grid;
  Grid = 0;
}

double C3DField::GetRealVal(CVector& p) {
  CVector pos = (p - (Min))/Delta; 
  return GetR(pos.x, pos.y, pos.z);
}

int C3DField::GetParticlePosition(CVector& p) {
  CVector pos = (p - (Min))/Delta; 
  return (int) (pos.x + pos.y*Width + pos.z*Width*Width);
}

CVector C3DField::GetParticlePositionInGrid(CVector& p) {
  return (p + (Min))/Delta; 
}

complex<double>& C3DField::Get(int x, int y, int z) {
  x = x % Width;
  y = y % Width;
  z = z % Width;
  return Grid[ x + y*Width + z*Width*Width ];
}


double& C3DField::GetR(int x, int y, int z) {
  return Get(x,y,z).real();
}

double& C3DField::GetRBoundary(const int& xx, const int& yy, const int& zz) {
  int x = xx % Width;
  int y = yy % Width;
  int z = zz % Width;
  return Get(x,y,z).real();
}


double& C3DField::GetI(int x, int y, int z) {
  return Get(x,y,z).imag();
}


void C3DField::Set(int x, int y, int z, double val){
  x = x % Width;
  y = y % Width;
  z = z % Width;

  Grid[ x + y*Width + z*Width*Width] = val;

}
void C3DField::Set(int x, int y, int z, complex<double> val){
  x = x % Width;
  y = y % Width;
  z = z % Width;

  Grid[ x + y*Width + z*Width*Width] = val;

}

void C3DField::Clear() {
  for (int i=0;i<Width*Width*Width;i++)
    Grid[i] = 0;
}

void C3DField::Clear(double d) {
  for (int i=0;i<Width*Width*Width;i++)
    Grid[i] = d;
}

void C3DField::toFloatGrid() {
  float x = 0;
  for (int i=0;i<Width*Width*Width;i++) {
    float scale = 1.0;
    float f = (Grid[i].real())*scale;
    /* x++;
     if (x>=Width) x = 0;
     f=x/Width;*/
    //f = x/(float)Width;
    dGrid[3*i] = f;
    dGrid[3*i+1] = f;
    dGrid[3*i+2] = f;
    //dGrid[i] = f;
   
  }
}

void C3DField::Initialize(int width, CVector min, CVector max, double boxsize) {

  BoxSize = boxsize;
  if (Grid && width==Width) {
    Clear();

  }
  else {
    Width = width;
    
    Delete();
    Grid = new complex<double>[Width*Width*Width]; 
    dGrid = new float[3*Width*Width*Width];
    Clear();
  }
  Min = min;
  Max = max;
  Delta = (Max-Min)*1.0/(double)Width;
}

void C3DField::InsertDirect(CVector& p, double val) {
  
  CVector pos = (p + (Min))/Delta; 
  Get(pos.x, pos.y, pos.z) +=val;
}

void C3DField::Differentiate(C3DField& Fx, C3DField& Fy, C3DField& Fz, double cutoff) {
  // Note: The current field should be in FOURIER SPACE!

  Fx.Initialize(Width, Min, Max,Fx.BoxSize);
  Fy.Initialize(Width, Min, Max,Fx.BoxSize);
  Fz.Initialize(Width, Min, Max,Fx.BoxSize);
  
  int Ng = Width;
  int w = 0;
  for (int l=w;l<Ng;l++)  {
    for (int m=w;m<Ng;m++) {
      for (int n=w;n<Ng;n++) {
	complex<double> rho = Get(l,m,n);
	complex<double> i(0,1);

	
	double ll = l;
	double mm = m;
	double nn = n;
	if (ll>Ng/2) ll-=Ng;
	if (mm>Ng/2) mm-=Ng;
	if (nn>Ng/2) nn-=Ng;
	
	if ((l==0 && m==0 && n==0) || (ll==Ng/2 || mm == Ng/2 || nn==Ng/2))  {
      	  
	}
	else
	  {
	    ll/=(double)Ng*2.0*CMath::pi;
	    mm/=(double)Ng*2.0*CMath::pi;
	    nn/=(double)Ng*2.0*CMath::pi;
	    double k2 = (ll*ll + mm*mm + nn*nn);
	    double C = exp(-k2*cutoff*cutoff);
	    Fx.Set(l,m,n, i*complex<double>((double)ll/k2,0)*rho*C);
	    Fy.Set(l,m,n, i*complex<double>((double)mm/k2,0)*rho*C);
	    Fz.Set(l,m,n, i*complex<double>((double)nn/k2,0)*rho*C);
	    
	  }
      } 
    }
  }

  Fx.FFTInv(Fx);
  Fy.FFTInv(Fy);
  Fz.FFTInv(Fz);
}




void C3DField::Correlation2(string Filename, CGraph& g) {

  /*  CGraph Correlation, CorrelationCount;
  CGraph w;
  w.Initialize(g.Bins);


  Correlation.Initialize(w.Bins);
  CorrelationCount.Initialize(w.Bins);
    
  double max = 1.5*Width;
  double min = 0.0;
  double delta = Width/(double)w.Bins;

  double l = Width/2.5;

  for (int x=0;x<Width;x++) {
    for (int y=0;y<Width;y++) {
      for (int z=0;z<Width;z++) {
	


	for (int i=0;i<Width;i++) {
	  for (int j=0;j<Width;j++) {
	    for (int k=0;k<Width;k++) {
	      
	      double d = (CVector(x,y,z) - CVector(i,j,k)).Length();
	      //if (d<Width)
		{
	      double val = GetR(x,y,z)*GetR(i,j,k);
	      int l = (int)(d/(double)delta);
	      //cout << d << " " << val << endl;
	      
	      if (l>=0 && l<w.Bins) {
		Correlation[l] += val; 
		CorrelationCount[l] += 1.0; 
	      }
	      }

	    }
	  }
	}
	
      }
    }
  }

  ofstream f(Filename.c_str(), ios::out );
  max = -1E10;
  for (int i=1;i<w.Bins;i++) {
    //f << i << " " << CorrelationCount[i] << endl;
    if (CorrelationCount[i]==0) CorrelationCount[i]=1;
    f << i << " " << Correlation[i]/(double)CorrelationCount[i] << endl;
    w[i] =  Correlation[i]/(double)CorrelationCount[i];
    if (w[i]>max) max=w[i];
  }
  f.close();
  CGraph h;

  max = -1E10;

  max = -1E10;
  w.FFTW(h);


  g.Initialize(w.Bins);
  for (int i=0;i<g.Bins;i++) 
    g[i] = h[i];
  
  for (int i=1;i<g.Bins;i++) {
    if (g[i]>max) max=g[i];
  }
  for (int i=0;i<g.Bins;i++)
    g[i]=g[i]/max;
  
  */
}


void C3DField::PowerSpectrum(CGraph& result, CGraph& Growth, int idx) {


    int Bins = Width;
    int Ng = Width;

    int maxb = 0.45*sqrt(3.0*(Bins/2)*(Bins/2));

    result.Initialize(maxb);
    for (int i=0;i<result.Bins; i++) 
      result[i] = 0.0;
    CGraph count;
    count.Initialize(maxb);

    double avg = 0.0;

      int v = 0;
      for (int l=v;l<Ng-v;l++)  {
	for (int m=v;m<Ng-v;m++) {
	  for (int n=v;n<Ng-v;n++) {

	    complex<double> d = Get(l,m,n);

	    int t =1;
	    if ((l == Ng/2 || m==Ng/2 || n == Ng/2 || (l*l+m*m+n*n)==0) || l==0 || m ==0 || n==0 || l==Ng-1 || m==Ng-1 || n == Ng-1) {
	      //cout << "eh " << x<< endl;
	      }
	    else 
	      {

		double xx = min(l, Ng-l);
		double yy = min(m, Ng-m);
		double zz = min(n, Ng-n);
		//if (xx>Ng/2) xx-=Ng;
		//if (yy>Ng/2) yy-=Ng;
		//if (zz>Ng/2) zz-=Ng;


		int k = 10.0*log10(sqrt(xx*xx + yy*yy +zz*zz) + 0.5);
		int kk =sqrt(xx*xx + yy*yy +zz*zz);
		double val = (norm(d));//*(k*k*k);
		if (kk>=3 && kk<=6) {
		  avg+=(val*val);
		}
		//int k = sqrt(xx*xx + yy*yy +zz*zz);
	    
		
		//cout << k  <<  ": " << Bins<< endl;
		if (k>0 && k<maxb) {
		  result[k]+=(val*val);//*norm(d);
		  count[k]+=1;
		}
		//else cout << k << endl;
	    }
	    

	  }
	  
	}
      }

      
    double max = -1E30;
    double min = 1E30;
    avg/=4.0;
    Growth[idx] = log(avg/1.0E11);
    //cout << avg << endl;
    result[0] = 0;
    for (int i=1;i<result.Bins;i++) {
      if (count[i]==0.0) 
	result[i] = 0;
      else
	result[i]/=(double)count[i];
      if (result[i]!=0)
	result[i] = log (result[i]);
      
      if (result[i]>max) max=result[i];
      if (result[i]<min) min=result[i];
      avg+=result[i];
    }
    //max = 1E8;
    for (int i=0;i<result.Bins;i++) 
      result[i] = (result[i] - min)/(2.0*max);//(avg/20.0);
    
    
 
}


void C3DField::SetupFFT(fftw_complex **in, fftw_complex** out, double scale) {
 
  double B = Width*Width*Width;
  if (!*out) *out = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * B);
  if (!*in) *in = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * B);

  for (int i=0;i<B;i++)  {
    complex<double> d = Grid[i]*scale; 
    (*in)[i][0] = d.real();
    (*in)[i][1] = d.imag();
  }
  
  for (int i=0;i<B;i++) {
    (*out)[i][0] = 0;
    (*out)[i][1] = 0;
  }

}

void C3DField::ResolveFFT(fftw_complex* out, C3DField* o) {
  for (int i=0;i<Width*Width*Width;i++) 
    o->Grid[i] = complex<double>(out[i][0], out[i][1]);
  
}



void C3DField::FFT(C3DField& o) {

    SetupFFT(&in, &out, 1);
    int B = Width*Width*Width;

    if (!PlanFFTInitialized) {
      PlanFFT = fftw_plan_dft_3d(Width, Width, Width,
				 in, out, FFTW_FORWARD,
				 FFTW_MEASURE);
      SetupFFT(&in, &out, 1);
      PlanFFTInitialized = true;
    }

    fftw_execute(PlanFFT); /* repeat as needed */

    o.Initialize(Width, Min, Max,BoxSize);

    ResolveFFT(out, &o);

    
}


void C3DField::FFTInv(C3DField& o) {
    
    double scale = 1.0/(double)(Width*Width*Width);


    SetupFFT(&in, &out, scale);
    
    if (!PlanFFTInvInitialized) {
      PlanFFTInv = fftw_plan_dft_3d(Width, Width, Width,
			 in, out, FFTW_BACKWARD,
			     FFTW_MEASURE);
      SetupFFT(&in, &out, scale);
      PlanFFTInvInitialized = true;
    }


    fftw_execute(PlanFFTInv); /* repeat as needed */
    
    o.Initialize(Width,Min,Max, BoxSize);
 
    ResolveFFT(out, &o);
    
}



void C3DField::Gaussian(double B) {
  int x[8] = {0,1,0,1,0,1,0,1};
  int y[8] = {0,0,1,1,0,0,1,1};
  int z[8] = {0,0,0,0,1,1,1,1};

  double sum = 0;
  for (int i=0;i<Width; i++)
    for (int j=0;j<Width; j++)
      for (int k=0;k<Width; k++) {
	double dist = 1E10;
	double val = 0;
	for (int l=0;l<8;l++) {
	  double d = (CVector(x[l],y[l],z[l])*Width - CVector(i,j,k)).Length();
	  d/=Width;
	  dist = min(dist,d);
	  val += exp(-dist*dist/(2.0*B*B));

	}
	
	sum+=val;
	Set(i,j,k,val);
      }
  //  sum/=2;
  for (int i=0;i<Width*Width*Width;i++)
    Grid[i]/=sum;

}

void C3DField::Mul(C3DField &other) {
  for (int i=0;i<Width*Width*Width; i++) {
    Grid[i]*=other.Grid[i];
  } 

}

void C3DField::Scale(double s) {
  for (int i=0;i<Width*Width*Width; i++) {
    Grid[i]*=s;
  } 
}

void C3DField::PhysicalScale(double FieldScale, double rho_c) {
  //  cout << Grid[255] << endl;

  for (int i=0;i<Width*Width*Width; i++) {
    Grid[i]= FieldScale*(Grid[i]/rho_c -1.0);
  } 
}


void C3DField::Copy(C3DField& s) {
  for (int i=0;i<Width*Width*Width; i++) {
    Grid[i] = s.Grid[i];
  } 
  
}


void C3DField::GenerateField(double scale, double D) {
  Clear();
  complex<double> I(0,1);
  for (int i=0;i<Width;i++) 
    for (int j=0;j<Width;j++) 
      for (int k=0;k<Width;k++)  {
	double A = CMath::RandomUniform();
	double B = CMath::RandomUniform();
	int ii = min(i,Width-i);
	int jj = min(j,Width-j);
	int kk = min(k,Width-k);
	
	double Kval = sqrt(ii*ii + jj*jj + kk*kk);
	if (Kval<0.1) Kval =0.1;
	double m = (Kval - scale);
	double Pk = D*1.0E5*exp(-m*m);
	complex<double> d = sqrt(-2.0*Pk*log(A))*exp(I*complex<double>(2.0*CMath::pi*B,0));
	Set(i,j,k,d);
      }
  

}
